import numpy as np
from mmcv.transforms import BaseTransform
from scipy.ndimage import map_coordinates
from scipy.spatial.transform import Rotation as R
import matplotlib.pyplot as plt



class RandomRotate3D(BaseTransform):
    def __init__(self,
                 degree: float,
                 prob: float = 0.5,
                 keys: list[str] = ["img", "gt_seg_map"]):
        self.degree = degree
        self.prob = prob
        self.keys = keys
        # 预计算最大旋转角的余弦值
        self.cos_theta = np.cos(np.deg2rad(degree))

    def _sample_rotation_matrix(self):
        axis = np.random.randn(3)
        axis /= np.linalg.norm(axis)
        angle = np.random.uniform(-self.degree, self.degree)
        return R.from_rotvec(np.deg2rad(angle) * axis).as_matrix()

    def _rotate_volume(self, array: np.ndarray, rot: np.ndarray):
        z, y, x = array.shape
        center = np.array([z/2, y/2, x/2])

        dz, dy, dx = np.indices((z, y, x))
        coords = np.stack([dz, dy, dx], axis=0).reshape(3, -1).astype(np.float32)

        coords_centered = coords.T - center
        coords_rotated = (rot @ coords_centered.T).T + center

        rotated = map_coordinates(
            array,
            [coords_rotated[:, 0], coords_rotated[:, 1], coords_rotated[:, 2]],
            order=1,
            mode='nearest'
        ).reshape(z, y, x)
        return rotated

    def _compute_minimal_bounds(self, shape):
        """计算理论最小有效区域"""
        z, y, x = shape
        center = np.array([z/2, y/2, x/2])
        
        # 计算对角线长度（从中心到角点的最大距离）
        max_radius = np.sqrt((z/2)**2 + (y/2)**2 + (x/2)**2)
        
        # 任意旋转后，点到中心的距离不变
        # 但在各轴上的投影最大可能偏移为 r*sin(theta)
        max_offset = max_radius * np.sqrt(1 - self.cos_theta**2)
        
        # 计算安全边界
        padding = np.ceil(max_offset).astype(int)
        
        return (padding, z-padding-1,
                padding, y-padding-1,
                padding, x-padding-1)

    def _compute_valid_crop_bounds(self, shape, rot):
        """计算旋转后有效区域 (zmin,zmax,ymin,ymax,xmin,xmax)"""
        z, y, x = shape
        center = np.array([z/2, y/2, x/2])

        # 8个角点
        corners = np.array([
            [0, 0, 0], [0, 0, x-1], [0, y-1, 0], [0, y-1, x-1],
            [z-1, 0, 0], [z-1, 0, x-1], [z-1, y-1, 0], [z-1, y-1, x-1]
        ], dtype=np.float32)

        corners_centered = corners - center
        corners_rotated = (rot @ corners_centered.T).T + center

        zmin, ymin, xmin = corners_rotated.min(axis=0)
        zmax, ymax, xmax = corners_rotated.max(axis=0)

        # 与原体素范围 [0, z-1], [0, y-1], [0, x-1] 做交集
        zmin, zmax = max(0, np.ceil(zmin)), min(z-1, np.floor(zmax))
        ymin, ymax = max(0, np.ceil(ymin)), min(y-1, np.floor(ymax))
        xmin, xmax = max(0, np.ceil(xmin)), min(x-1, np.floor(xmax))

        return int(zmin), int(zmax), int(ymin), int(ymax), int(xmin), int(xmax)

    def _center_crop(self, array: np.ndarray, bounds):
        zmin, zmax, ymin, ymax, xmin, xmax = bounds
        return array[zmin:zmax+1, ymin:ymax+1, xmin:xmax+1]

    def transform(self, results):
        if np.random.rand() < self.prob:
            # 1. 生成随机旋转矩阵
            rot = self._sample_rotation_matrix()
            
            # 2. 计算理论最小边界（基于最大旋转角）
            bounds = self._compute_minimal_bounds(
                results[self.keys[0]].shape
            )
            
            # 3. 先旋转再裁剪
            for key in self.keys:
                # 旋转整个体素
                rotated = self._rotate_volume(results[key], rot)
                # 在旋转后的体素上裁剪安全区域
                results[key] = self._center_crop(rotated, bounds)
        
        return results


def create_test_cube(size=64):
    """创建带有立方体特征的测试体积"""
    volume = np.zeros((size, size, size))
    center = size // 2
    cube_size = size // 4
    
    # 创建中心立方体
    start = center - cube_size // 2
    end = center + cube_size // 2
    volume[start:end, start:end, start:end] = 1
    
    # 添加边缘标记点
    volume[0:3, 0:3, 0:3] = 2  # 角点标记
    volume[-3:, -3:, -3:] = 2  # 对角角点标记
    
    return volume


def plot_volume_slices(volume, title=""):
    """简单地显示三个正交切片"""
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    mid = [s//2 for s in volume.shape]
    
    imgs = [
        axes[0].imshow(volume[mid[0],:,:], vmin=0, vmax=2),
        axes[1].imshow(volume[:,mid[1],:], vmin=0, vmax=2),
        axes[2].imshow(volume[:,:,mid[2]], vmin=0, vmax=2)
    ]
    
    titles = ['XY平面', 'XZ平面', 'YZ平面']
    for ax, title in zip(axes, titles):
        ax.set_title(title)
    
    fig.colorbar(imgs[0], ax=axes.ravel().tolist())
    plt.suptitle(title)
    return fig


def test_random_rotate():
    # 创建测试数据
    volume = create_test_cube(64)
    
    # 创建变换实例
    transform = RandomRotate3D(
        degree=45,
        prob=1.0,
        keys=['img']
    )
    
    # 应用变换
    results = {'img': volume}
    results = transform.transform(results)
    rotated_volume = results['img']
    
    # 显示原始体积
    print(volume.shape)
    # plot_volume_slices(volume, "原始体积")
    # plt.show()
    
    # 显示旋转后的体积
    print(rotated_volume.shape)
    plot_volume_slices(rotated_volume, "旋转后的体积")
    plt.show()



if __name__ == "__main__":
    test_random_rotate()